#encoding=utf-8
import numpy as np
import cv2
from data.bbox.bbox_dataset import DetectionDataset
class InstancesFormula(DetectionDataset):
    def parser_xml(self, xml_path):
        import xml.etree.ElementTree as ET
        import logging
        oneimg = {}
        oneimg['bboxes'] = []
        oneimg["latex"] = []
        oneimg["latex_normalized"] = []
        oneimg["seg"] = []
        try:
            dom = ET.parse(xml_path)
        except Exception as e:
            logging.warning("{}_{}".format(e, xml_path))
            return None
        root = dom.getroot()
        img_name = root.findall('filename')[0].text
        oneimg['img_path'] = xml_path[:-4] + u".png"
        for objects in root.findall('object'):
            name = objects.find('name').text
            points = list(objects.find('polygon'))
            if len(points) != 4:
                logging.warning("find illegal label in file:{}.xml. ".format(img_name))
                return None
            toxy = lambda x: [int(x[0]), int(x[1])]
            points = map(lambda x: toxy(x.text.strip().split(",")), points)
            points = list(points)
            points = np.array(points)
            xmin = np.min(points[:, 0])
            ymin = np.min(points[:, 1])
            xmax = np.max(points[:, 0])
            ymax = np.max(points[:, 1])
            if name.startswith(u"###"):  # type: str
                pass
            elif name.startswith(u"@@@"):
                pass
            else:
                oneimg["seg"].append(points[np.newaxis])
                oneimg["bboxes"].append([xmin,ymin,xmax,ymax,0])
        self.classes = ["chars",]
        return oneimg
    def __init__(self, xml_root = u"/data3/zyx/ocr/For_Detection/A_Classs"):
        super(InstancesFormula,self).__init__()
        from utils.common import lsdir
        xmls = lsdir(xml_root,suffix=u".xml")
        objs = [self.parser_xml(xml) for xml  in xmls]
        self.objs = objs
    def __getitem__(self, idx):
        oneimg = self.objs[idx]
        img_pah = oneimg['img_path']
        bbox = oneimg['bboxes']
        seg = oneimg["seg"]
        return cv2.imread(img_pah), bbox, seg
    def at_with_image_path(self, idx):
        oneimg = self.objs[idx]
        img_pah = oneimg['img_path']
        bbox = oneimg['bboxes']
        seg = oneimg["seg"]
        return img_pah.encode(encoding="utf-8"), np.array(bbox),seg

    def __len__(self):
        return len(self.objs)
da = InstancesFormula().viz_seg()